import streamlit as st
import pandas as pd
from zhipuai import ZhipuAI
import json
from streamlit_echarts import st_echarts

# === 初始化 ===
# 在文件开头保留唯一的set_page_config调用
st.set_page_config(layout="wide", page_title="智聘通")

# === 数据加载 ===
@st.cache_data
def load_data():
    data = pd.read_csv('拉勾网2023招聘数据.csv')
    
    # 修复后的薪资解析逻辑
    def parse_salary(s):
        if isinstance(s, str):
            s = s.replace('k', '').replace('K', '')  # 处理大小写K
            parts = s.split('-')
            try:
                return [float(parts[0]), float(parts[1]) if len(parts)>1 else float(parts[0])]
            except:
                return [0, 0]
        return [0, 0]
    
    salary_data = data['salary'].apply(parse_salary)
    data['min_salary'] = salary_data.apply(lambda x: x[0])
    data['max_salary'] = salary_data.apply(lambda x: x[1])
    data['avg_salary'] = (data['min_salary'] + data['max_salary']) / 2 * 1000
    data['salary_mid'] = data['salary'].apply(
        lambda s: (float(s.split('-')[0].replace('k','').replace('K','')) + 
                  float(s.split('-')[1].replace('k','').replace('K','')))/2*1000 
        if '-' in s else None
    )
    
    # 技能标签处理
    data['skillLables'] = data['skillLables'].apply(
        lambda x: eval(x) if isinstance(x, str) else []
    )
    # 新增薪资上下限字段（兼容实验2）
    data['salary_lower'] = data['min_salary']
    data['salary_upper'] = data['max_salary']
    return data

data = load_data()

# === 主界面设计 ===
tab1, tab2, tab3 = st.tabs(["📊 数据看板", "🔍 技能关联分析", "💬 AI咨询"])

with tab1:
    st.header("招聘市场数据看板")
    
    # 添加筛选控件
    col1, col2 = st.columns(2)
    with col1:
        selected_city = st.selectbox("选择城市", ["全部"] + list(data['city'].unique()))
    with col2:
        selected_position = st.selectbox("选择职位", ["全部"] + list(data['positionName'].unique()))
    
    # 添加图表类型选择器
    region_chart_type = st.sidebar.selectbox('地区图表类型', ['柱状图', '折线图', '饼图'])
    salary_chart_type = st.sidebar.selectbox('薪资图表类型', ['柱状图', '折线图', '散点图', '雷达图'])
    position_chart_type = st.sidebar.selectbox('岗位图表类型', ['柱状图', '折线图', '饼图', '散点图'])
    
    # 添加薪资范围选择器
    salary_range = st.sidebar.slider(
        '选择薪资范围(元)', 
        min_value=int(data['salary_mid'].min()), 
        max_value=int(data['salary_mid'].max()), 
        value=(int(data['salary_mid'].min()), int(data['salary_mid'].max()))
    )
    
    # 添加工作经验选择器
    work_experience = st.sidebar.selectbox('选择工作经验', ["全部"] + list(data['workYear'].unique()))
    
    # 添加学历要求选择器
    education_requirement = st.sidebar.selectbox('选择学历要求', ["全部"] + list(data['education'].unique()))
    
    # 1. 地区分布图表
    region_option = {
        "title": {"text": '地区分布'},
        "tooltip": {
            "trigger": 'item',
            "formatter": '{a}<br/>{b}: {c}个职位 ({d}%)'
        },
        "xAxis": {
            "type": 'category',
            "data": data['city'].value_counts().head(10).index.tolist()
        },
        "yAxis": {},
        "series": [
            {
                "name": '职位数量',
                "type": 'bar' if region_chart_type == '柱状图' else 'line' if region_chart_type == '折线图' else 'pie',
                "data": data['city'].value_counts().head(10).tolist(),
                "label": {
                    "show": True,
                    "formatter": '{b}: {c} ',
                    "position": 'outside',
                    "fontSize": 12,
                    "fontWeight": 'bold',
                    "emphasis": {
                        "show": True,
                        "fontSize": 14,
                        "fontWeight": 'bold'
                    }
                },
                "emphasis": {
                    "itemStyle": {
                        "shadowBlur": 10,
                        "shadowOffsetX": 0,
                        "shadowColor": 'rgba(0, 0, 0, 0.5)'
                    },
                    "label": {
                        "show": True
                    }
                }
            }
        ]
    }

    if region_chart_type == '饼图':
        region_option['series'][0]['radius'] = '50%'
        region_option['series'][0]['center'] = ['50%', '50%']

    # 2. 薪资分布图表
    salary_option = {
        "title": {
            "text": '薪资分布',
            "subtext": '数据来自拉勾网',
            "left": 'center'
        },
        "tooltip": {
            "trigger": 'axis',
            "axisPointer": {
                "type": 'cross',
                "crossStyle": {
                    "color": '#999'
                }
            }
        },
        "legend": {
            "data": ['岗位数量', '平均薪资'],
            "left": 'right'
        },
        "xAxis": {
            "type": 'category',
            "data": ['<10k', '10-20k', '20-30k', '30-40k', '>40k'],
            "axisPointer": {
                "type": 'shadow'
            }
        },
        "yAxis": [
            {
                "type": 'value',
                "name": '岗位数量',
                "min": 0,
                "axisLabel": {
                    "formatter": '{value}'
                }
            },
            {
                "type": 'value',
                "name": '平均薪资',
                "min": 0,
                "axisLabel": {
                    "formatter": '{value}元'
                }
            }
        ],
        "series": [
            {
                "name": '岗位数量',
                "type": 'line' if salary_chart_type == '折线图' else 'bar' if salary_chart_type == '柱状图' else 'pie',
                "data": [
                    len(data[data['salary_mid'] < 10000]),
                    len(data[(data['salary_mid'] >= 10000) & (data['salary_mid'] < 20000)]),
                    len(data[(data['salary_mid'] >= 20000) & (data['salary_mid'] < 30000)]),
                    len(data[(data['salary_mid'] >= 30000) & (data['salary_mid'] < 40000)]),
                    len(data[data['salary_mid'] >= 40000])
                ],
                "itemStyle": {
                    "color": '#5470C6'
                }
            },
            {
                "name": '平均薪资',
                "type": 'line' if salary_chart_type == '折线图' else 'bar' if salary_chart_type == '柱状图' else 'pie',
                "yAxisIndex": 1,
                "data": [
                    data[data['salary_mid'] < 10000]['salary_mid'].mean(),
                    data[(data['salary_mid'] >= 10000) & (data['salary_mid'] < 20000)]['salary_mid'].mean(),
                    data[(data['salary_mid'] >= 20000) & (data['salary_mid'] < 30000)]['salary_mid'].mean(),
                    data[(data['salary_mid'] >= 30000) & (data['salary_mid'] < 40000)]['salary_mid'].mean(),
                    data[data['salary_mid'] >= 40000]['salary_mid'].mean()
                ],
                "itemStyle": {
                    "color": '#EE6666'
                }
            }
        ]
    }

    if salary_chart_type == '雷达图':
        radar_data = [
            {
                'name': '岗位数量',
                'value': [
                    len(data[data['salary_mid'] < 10000]),
                    len(data[(data['salary_mid'] >= 10000) & (data['salary_mid'] < 20000)]),
                    len(data[(data['salary_mid'] >= 20000) & (data['salary_mid'] < 30000)]),
                    len(data[(data['salary_mid'] >= 30000) & (data['salary_mid'] < 40000)]),
                    len(data[data['salary_mid'] >= 40000])
                ],
                'itemStyle': {'color': '#5470C6'}
            },
            {
                'name': '平均薪资',
                'value': [
                    data[data['salary_mid'] < 10000]['salary_mid'].mean(),
                    data[(data['salary_mid'] >= 10000) & (data['salary_mid'] < 20000)]['salary_mid'].mean(),
                    data[(data['salary_mid'] >= 20000) & (data['salary_mid'] < 30000)]['salary_mid'].mean(),
                    data[(data['salary_mid'] >= 30000) & (data['salary_mid'] < 40000)]['salary_mid'].mean(),
                    data[data['salary_mid'] >= 40000]['salary_mid'].mean()
                ],
                'itemStyle': {'color': '#EE6666'}
            }
        ]
        
        salary_option = {
            'title': {'text': '薪资分布雷达图', 'left': 'center'},
            'tooltip': {'trigger': 'item'},
            'legend': {'data': ['岗位数量', '平均薪资'], 'left': 'right'},
            'radar': {
                'indicator': [
                    {'name': '<10k', 'max': max([d['value'][0] for d in radar_data]) * 1.2},
                    {'name': '10-20k', 'max': max([d['value'][1] for d in radar_data]) * 1.2},
                    {'name': '20-30k', 'max': max([d['value'][2] for d in radar_data]) * 1.2},
                    {'name': '30-40k', 'max': max([d['value'][3] for d in radar_data]) * 1.2},
                    {'name': '>40k', 'max': max([d['value'][4] for d in radar_data]) * 1.2}
                ]
            },
            'series': [{
                'name': '薪资分布',
                'type': 'radar',
                'data': radar_data
            }]
        }

    # 3. 学历和工作经验分布图表
    position_option = {
        "title": {"text": '学历和工作经验分布'},
        "tooltip": {},
        "legend": {
            "data": ['学历分布', '工作经验分布'],
            "top": 'bottom'
        },
        "xAxis": {
            "type": 'category',
            "data": data['education'].value_counts().index.tolist()
        },
        "yAxis": {
            "type": 'value'
        },
        "series": [
            {
                "name": '学历分布',
                "type": 'bar' if position_chart_type == '柱状图' else 'line' if position_chart_type == '折线图' else 'pie',
                "data": data['education'].value_counts().tolist(),
                "itemStyle": {
                    "color": '#5470C6'
                }
            },
            {
                "name": '工作经验分布',
                "type": 'line' if position_chart_type == '折线图' else 'bar' if position_chart_type == '柱状图' else 'pie',
                "data": data['workYear'].value_counts().tolist(),
                "itemStyle": {
                    "color": '#EE6666'
                }
            }
        ]
    }

    if position_chart_type == '饼图':
        position_option['series'][0]['radius'] = '50%'
        position_option['series'][0]['center'] = ['50%', '50%']
        position_option['series'][1]['radius'] = '50%'
        position_option['series'][1]['center'] = ['50%', '50%']

    # 根据选择过滤数据
    if selected_city != '全部' and selected_position != '全部':
        filtered_data = data[(data['city'] == selected_city) & (data['positionName'] == selected_position)]
    elif selected_city != '全部':
        filtered_data = data[data['city'] == selected_city]
    elif selected_position != '全部':
        filtered_data = data[data['positionName'] == selected_position]
    else:
        filtered_data = data

    # 应用薪资范围过滤
    filtered_data = filtered_data[
        (filtered_data['salary_mid'] >= salary_range[0]) & 
        (filtered_data['salary_mid'] <= salary_range[1])
    ]

    # 应用工作经验过滤
    if work_experience != '全部':
        filtered_data = filtered_data[filtered_data['workYear'] == work_experience]
        
    # 应用学历要求过滤
    if education_requirement != '全部':
        filtered_data = filtered_data[filtered_data['education'] == education_requirement]

    # 更新图表数据
    region_option['xAxis']['data'] = filtered_data['city'].value_counts().head(10).index.tolist()
    region_option['series'][0]['data'] = filtered_data['city'].value_counts().head(10).tolist()

    # 显示图表
    col1, col2 = st.columns(2)
    with col1:
        st_echarts(options=region_option, height="400px")
    with col2:
        st_echarts(options=salary_option, height="400px")
    
    st_echarts(options=position_option, height="400px")

    # 显示过滤后的数据
    st.dataframe(filtered_data)

with tab2:
    # 整合Untitled.ipynb的技能关联分析
    st.header("岗位技能关联分析")
    
    # 岗位→技能查询
    st.subheader("岗位技能要求分析")
    position = st.selectbox("选择岗位", data['positionName'].unique())
    position_skills = data[data['positionName'] == position]['skillLables'].explode().value_counts()
    st_echarts({
        "tooltip": {
            "trigger": "item",
            "formatter": "{b}: {c} 个"
        },
        "xAxis": {"type": "category", "data": position_skills.head(10).index.tolist()},
        "yAxis": {"type": "value"},
        "series": [{"data": position_skills.head(10).tolist(), "type": "bar"}]
    })
    
    # 技能→岗位查询
    st.subheader("技能对应岗位分析")
    skill = st.selectbox("选择技能", pd.Series(sum(data['skillLables'].tolist(), [])).unique())
    skill_positions = data[data['skillLables'].apply(lambda x: skill in x)]['positionName'].value_counts()
    st_echarts({
        "tooltip": {
            "trigger": "item",
            "formatter": "{b}: {c} ({d}%)"
        },
        "series": [{
            "type": "pie",
            "data": [{"value": v, "name": k} for k,v in skill_positions.items()],
            "radius": "50%"
        }]
    })

with tab3:
    # 整合LLM函数调用实验的AI对话功能
    if 'history' not in st.session_state:
        st.session_state.history = []
    
    # 显示对话历史
    for msg in st.session_state.history:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])
            if "chart" in msg:
                st_echarts(msg["chart"])
        
        try:
            response = client.chat.completions.create(
                model="glm-4",
                messages=[{"role": m["role"], "content": m["content"]} for m in st.session_state.history],
                tools=tools,
                temperature=0.7
            )
    
            ai_msg = response.choices[0].message
            if ai_msg.tool_calls:
                tool_results = handle_tool_calls(ai_msg.tool_calls)
                st.session_state.history.extend(tool_results)
        
        # 处理工具调用结果
            for result in tool_results:
                try:
                    result_data = pd.read_json(result['content'])
                    tool_name = result['name']  # 获取工具名称
                    
                    if tool_name == "query_work_experience_requirement":
                       # 解析工作经验数据
                        work_exp_data = pd.DataFrame(result_data['work_experience'].iloc[0])
                        # 生成柱状图配置
                        chart_option = {
                            "title": {"text": f"{result_data['city'].iloc[0]}地区{result_data['job_name'].iloc[0]}工作经验要求分布"},
                            "tooltip": {
                                "trigger": "item",
                                "formatter": "{b}: {c} 个岗位"
                            },
                            "xAxis": {
                                "type": "category",
                                "data": work_exp_data['work_year'].tolist()
                            },
                            "yAxis": {
                                "type": "value",
                                "name": "岗位数量"
                            },
                            "series": [{
                                "type": "bar",
                                "data": work_exp_data['count'].tolist(),
                                "itemStyle": {
                                    "color": "#5470C6",
                                    "barBorderRadius": 5
                                },
                                "label": {"show": True, "formatter": "{c}"}
                            }]
                        }
                        # 添加到对话历史
                        st.session_state.history.append({
                            "role": "assistant",
                            "content": f"{result_data['city'].iloc[0]}地区{result_data['job_name'].iloc[0]}的工作经验要求分布如下：",
                            "chart": chart_option
                        })
                    elif not result_data.empty:
                        # 保留原有其他工具的图表逻辑（如岗位数量柱状图）
                        st.session_state.history.append({
                            "role": "assistant",
                            "content": f"找到{len(result_data)}条相关职位信息",
                            "chart": {
                                "xAxis": {"type": "category", "data": result_data['positionName'].tolist()},
                                "yAxis": {"type": "value"},
                                "series": [{"data": result_data['avg_salary'].tolist(), "type": "bar"}]
                            }
                        })
                except Exception as e:
                    st.session_state.history.append({
                        "role": "assistant",
                        "content": f"处理工具结果时出错: {str(e)}"
                    })
                    
                    st.rerun()
        except Exception as e:
            st.session_state.history.append({
                "role": "assistant",
                "content": f"请求 AI 回复时出错: {str(e)}"
            })

# === 工具函数定义 ===
tools = [
    {
        "type": "function",
        "function": {
            "name": "query_language_jobs",
            "description": "查询指定编程语言相关岗位的详细信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "language": {"type": "string", "description": "编程语言名称"},
                    "top_n": {"type": "number", "description": "返回前N个岗位（可选，默认10）"}
                },
                "required": ["language"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_average_salary",
            "description": "根据城市和职位名称查询该职位的平均薪资",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "目标城市，如'上海'"},
                    "job_name": {"type": "string", "description": "职位名称，如'算法工程师'"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_max_salary_job",
            "description": "查询某城市薪资最高的岗位",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "top_n": {"type": "number", "description": "返回前N个薪资最高岗位（可选，默认5）"}
                },
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_salary_range",
            "description": "查询某城市某职位的薪资区间",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "job_name": {"type": "string", "description": "职位名称"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_job_demand",
            "description": "查询某城市某职位的岗位需求量",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "job_name": {"type": "string", "description": "职位名称"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_language_related_jobs",
            "description": "查询掌握指定编程语言可以从事的岗位及其需求量",
            "parameters": {
                "type": "object",
                "properties": {
                    "language": {"type": "string", "description": "编程语言名称"},
                    "top_n": {"type": "number", "description": "返回前N个岗位（可选，默认5）"}
                },
                "required": ["language"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_language_learning_recommendations",
            "description": "查询指定编程语言相关的技术学习推荐",
            "parameters": {
                "type": "object",
                "properties": {
                    "language": {"type": "string", "description": "编程语言名称"},
                    "top_n": {"type": "number", "description": "返回前N个推荐（可选，默认10）"}
                },
                "required": ["language"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_cpp_learning_recommendations",
            "description": "查询C++相关的技术学习推荐",
            "parameters": {
                "type": "object",
                "properties": {
                    "top_n": {"type": "number", "description": "返回前N个推荐（可选，默认10）"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_highest_salary_job",
            "description": "查询指定城市薪资最高的岗位",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "top_n": {"type": "number", "description": "返回前N个薪资最高岗位(可选，默认5)"}
                },
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_experience_requirement",
            "description": "根据城市和职位名称查询该职位的常见工作经验要求",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "目标城市"},
                    "job_name": {"type": "string", "description": "职位名称"}
                },
                "required": ["city", "job_name"]
            }
        }
    }
]

# === 增强的AI回复处理逻辑 ===
def handle_tool_calls(tool_calls):
    results = []
    for tool_call in tool_calls:
        func_name = tool_call.function.name
        args = json.loads(tool_call.function.arguments)
        
        if func_name == "query_experience_requirement":
            city = args["city"]
            job_name = args["job_name"]
            # 从数据集中筛选对应城市和职位的记录，统计工作经验分布
            filtered = data[
                (data['city'] == city) & 
                (data['positionName'] == job_name)
            ]
            # 统计每个工作经验的岗位数量
            work_exp_counts = filtered['workYear'].value_counts().reset_index()
            work_exp_counts.columns = ["work_year", "count"]
            # 转换为JSON格式返回
            results.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": func_name,
                "content": json.dumps({
                    "city": city,
                    "job_name": job_name,
                    "work_experience": work_exp_counts.to_dict("records")
                })
            })
        elif func_name == "query_language_jobs":
            filtered = data[data['skillLables'].apply(lambda x: args["language"] in x)]
            results.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": func_name,
                "content": filtered.head(args.get("top_n", 10)).to_json(orient='records')
            })
        elif func_name == "query_average_salary":
            filtered = data[(data['city'] == args["city"]) & 
                          (data['positionName'] == args["job_name"])]
            avg_salary = filtered['avg_salary'].mean()
            results.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": func_name,
                "content": json.dumps({"average_salary": avg_salary})
            })
        elif func_name == "query_max_salary_job":
            filtered = data[data['city'] == args["city"]]
            top_jobs = filtered.nlargest(args.get("top_n", 5), 'avg_salary')
            results.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": func_name,
                "content": top_jobs.to_json(orient='records')
            })
        elif func_name == "query_salary_range":
            city = args["city"]
            job_name = args["job_name"]
            filtered = data[(data['city'] == args["city"]) & 
                          (data['positionName'] == args["job_name"])]

            min_salary = filtered['min_salary'].min() * 1000  
            max_salary = filtered['max_salary'].max() * 1000  

            results.append({
                "name": func_name,
                "content": json.dumps({
                    "city": city,
                    "job_name": job_name,
                    "min": round(min_salary, 2),
                    "max": round(max_salary, 2)
                })
            })
            
        elif func_name == "query_job_demand":
            count = len(data[(data['city'] == args["city"]) & 
                           (data['positionName'] == args["job_name"])])
            results.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": func_name,
                "content": json.dumps({"count": count})
            })
        elif func_name == "query_language_related_jobs":
            language = args["language"]
            top_n = args.get("top_n", 5)
            
            # 从数据集中筛选包含该编程语言的岗位（假设技能标签字段为 skillLables）
            # 注意：确保数据集中 skillLables 字段已正确解析为列表（如 ["Python", "Java"]）
            filtered = data[
                data['skillLables'].apply(lambda skills: language in skills)
            ]
            
            # 统计岗位需求量（每个岗位出现的次数）
            job_demand = filtered['positionName'].value_counts().head(top_n)
            # 转换为 DataFrame 格式（用于后续可视化）
            job_demand_df = pd.DataFrame({
                "positionName": job_demand.index.tolist(),
                "demand": job_demand.tolist()
            })
            
            # 结果转换为 JSON 字符串（供 pd.read_json 解析）
            results.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": func_name,
                "content": job_demand_df.to_json(orient="records")
            })
        elif func_name == "query_language_learning_recommendations":
            language = args["language"]
            top_n = args.get("top_n", 10)
            # 这里可以添加根据语言返回不同学习推荐的逻辑
            recommendations = {
                "Python": ["基础语法", "Web开发", "数据分析", "机器学习", "自动化脚本"],
                "Java": ["核心语法", "Spring框架", "微服务", "JVM原理", "并发编程"],
                "C++": ["现代C++特性", "操作系统与系统编程", "高性能计算", "C++标准库", "开源项目贡献"],
                "JavaScript": ["ES6+特性", "Node.js", "前端框架", "TypeScript", "Web性能优化"]
            }
            return {
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": func_name,
                "content": json.dumps(recommendations.get(language, [])[:top_n])
            }
        elif func_name == "query_cpp_learning_recommendations":
            pass
        elif func_name == "query_work_experience_requirement":
            city = args["city"]
            job_name = args["job_name"]
            # 从数据集中筛选对应城市和职位的记录，统计工作经验分布
            filtered = data[
                (data['city'] == city) & 
                (data['positionName'] == job_name)
            ]
            # 统计每个工作经验的岗位数量
            work_exp_counts = filtered['workYear'].value_counts().reset_index()
            work_exp_counts.columns = ["work_year", "count"]
            # 转换为JSON格式返回（包含城市和职位信息）
            results.append({
                "name": func_name,
                "content": json.dumps({
                    "city": city,
                    "job_name": job_name,
                    "work_experience": work_exp_counts.to_dict("records")
                })
            })
    return results

# === 增强的图表生成函数 ===
def generate_chart(message):
    # 根据AI回复内容生成相应图表
    if "薪资" in message.content:
        return {
            "xAxis": {"type": "category", "data": ["最低", "平均", "最高"]},
            "yAxis": {"type": "value"},
            "series": [{
                "data": [
                    data['min_salary'].mean(),
                    data['avg_salary'].mean(),
                    data['max_salary'].mean()
                ],
                "type": "bar"
            }]
        }
    elif "地区" in message.content:
        region_counts = data['city'].value_counts().head(10)
        return {
            "xAxis": {"type": "category", "data": region_counts.index.tolist()},
            "yAxis": {"type": "value"},
            "series": [{"data": region_counts.tolist(), "type": "bar"}]
        }
    elif "技能" in message.content:
        # 新增技能关联图表
        skill = message.content.split("技能")[1].strip()
        if skill in pd.Series(sum(data['skillLables'].tolist(), [])).unique():
            skill_positions = data[data['skillLables'].apply(lambda x: skill in x)]['positionName'].value_counts()
            return {
                "series": [{
                    "type": "pie",
                    "data": [{"value": v, "name": k} for k,v in skill_positions.items()],
                    "radius": "50%"
                }]
            }
    
    # 新增编程语言相关岗位图表
    if "编程语言" in message.content or "语言" in message.content:
        language = next((word for word in message.content.split() 
                        if word in ["Python", "Java", "C++", "JavaScript"]), None)
        if language:
            language_jobs = data[data['skillLables'].apply(lambda x: language in x)]
            position_counts = language_jobs['positionName'].value_counts().head(10)
            return {
                "xAxis": {"type": "category", "data": position_counts.index.tolist()},
                "yAxis": {"type": "value"},
                "series": [{"data": position_counts.tolist(), "type": "bar"}],
                "title": {"text": f"{language}相关热门岗位"}
            }
    
    # 新增最高薪资岗位图表
    if "最高薪资" in message.content or "薪资最高" in message.content:
        city = next((word for word in message.content.split() 
                    if word in data['city'].unique()), None)
        if city:
            top_jobs = data[data['city'] == city].nlargest(5, 'avg_salary')
            return {
                "xAxis": {"type": "category", "data": top_jobs['positionName'].tolist()},
                "yAxis": {"type": "value"},
                "series": [{"data": top_jobs['avg_salary'].tolist(), "type": "bar"}],
                "title": {"text": f"{city}薪资最高岗位"}
            }
    elif "学习推荐" in message.content:
        return {
            "series": [{
                "type": "pie",
                "radius": ['30%', '70%'],
                "data": [
                    {"value": 9, "name": "现代C++特性"},
                    {"value": 8, "name": "操作系统与系统编程"},
                    {"value": 7, "name": "高性能计算"},
                    {"value": 6, "name": "C++标准库"},
                    {"value": 5, "name": "开源项目贡献"}
                ]
            }]
        }
    return None


# === 初始化智谱AI客户端 ===
def init_ai(api_key):
    if not api_key or len(api_key) < 10:
        raise ValueError("无效的API Key")
    return ZhipuAI(api_key=api_key)

# === 工具函数定义 ===
tools = [
    {
        "type": "function",
        "function": {
            "name": "query_language_jobs",
            "description": "查询指定编程语言相关岗位的详细信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "language": {"type": "string", "description": "编程语言名称"},
                    "top_n": {"type": "number", "description": "返回前N个岗位（可选，默认10）"}
                },
                "required": ["language"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_average_salary",
            "description": "根据城市和职位名称查询该职位的平均薪资",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "目标城市，如'上海'"},
                    "job_name": {"type": "string", "description": "职位名称，如'算法工程师'"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_max_salary_job",
            "description": "查询某城市薪资最高的岗位",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "top_n": {"type": "number", "description": "返回前N个薪资最高岗位（可选，默认5）"}
                },
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_salary_range",
            "description": "根据城市和职位名称查询该职位的薪资范围",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "目标城市"},
                    "job_name": {"type": "string", "description": "职位名称"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_experience_requirement",
            "description": "根据城市和职位名称查询该职位的常见工作经验要求",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "目标城市"},
                    "job_name": {"type": "string", "description": "职位名称"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_work_experience_requirement",
            "description": "查询某城市某职位的工作经验要求分布（不同工作经验对应的岗位数量）",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "目标城市，如'上海'"},
                    "job_name": {"type": "string", "description": "职位名称，如'Java工程师'"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_job_demand",
            "description": "查询某城市某职位的岗位需求量",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "job_name": {"type": "string", "description": "职位名称"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_language_related_jobs",
            "description": "查询掌握指定编程语言可以从事的岗位及其需求量",
            "parameters": {
                "type": "object",
                "properties": {
                    "language": {"type": "string", "description": "编程语言名称"},
                    "top_n": {"type": "number", "description": "返回前N个岗位（可选，默认5）"}
                },
                "required": ["language"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_language_learning_recommendations",
            "description": "查询指定编程语言相关的技术学习推荐",
            "parameters": {
                "type": "object",
                "properties": {
                    "language": {"type": "string", "description": "编程语言名称"},
                    "top_n": {"type": "number", "description": "返回前N个推荐（可选，默认10）"}
                },
                "required": ["language"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_cpp_learning_recommendations",
            "description": "查询C++相关的技术学习推荐",
            "parameters": {
                "type": "object",
                "properties": {
                    "top_n": {"type": "number", "description": "返回前N个推荐（可选，默认10）"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_highest_salary_job",
            "description": "查询指定城市薪资最高的岗位",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "top_n": {"type": "number", "description": "返回前N个薪资最高岗位(可选，默认5)"}
                },
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_top_jobs_in_city",
            "description": "查询指定城市中岗位数量最多的前N个职位",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "目标城市名称，如'上海'"},
                    "top_n": {"type": "number", "description": "返回前N个岗位（可选，默认5）"}
                },
                "required": ["city"]
            }
        }
    }
]

# 查询指定城市薪资最高的岗位
def query_max_salary_job(city, data, top_n=1):
    city_data = data[data['city'] == city].dropna(subset=['salary_mid'])
    if city_data.empty:
        return {}
    # 按薪资中位数降序排序，取第1个岗位
    top_jobs = city_data.groupby('positionName')['salary_mid'].mean() \
        .sort_values(ascending=False) \
        .head(top_n)  # 强制取第1个
    return top_jobs.to_dict()

# 查询指定城市的指定岗位的平均薪资
def query_average_salary(city, job_name, data):
    # 检查必要字段是否存在
    required_columns = ['city', 'positionName', 'salary_mid']
    if not all(col in data.columns for col in required_columns):
        return f"数据缺失必要字段：{[col for col in required_columns if col not in data.columns]}"

    # 筛选数据
    job_data = data[(data['city'] == city) & (data['positionName'] == job_name)]
    if job_data.empty:
        return {"error": f"未找到{city}的{job_name}岗位数据"}
    if job_data['salary_mid'].isna().all():
        return {"error": f"{city}的{job_name}岗位薪资数据缺失"}

    avg_salary = job_data['salary_mid'].mean()
    return {f"{city}{job_name}平均薪资": avg_salary}

def query_salary_range(city, job_name, data):
        job_data = data[(data['city'] == city) & (data['positionName'] == job_name)]
        if not job_data.empty:
            return (job_data['salary_lower'].min(), job_data['salary_upper'].max())
        else:
            return "未找到相关职位"

# === 自定义函数：查询地区职位工作经验要求 ===
def query_experience_requirement(city, job_name):
    job_data = data[(data['city'] == city) & (data['positionName'] == job_name)]
    if not job_data.empty:
        exp_req = job_data['workYear'].value_counts().to_dict()
        return [{"work_year": k, "count": v} for k, v in exp_req.items()]
    else:
        return "未找到相关职位"

    # 查询指定城市的指定岗位的岗位数量
def query_job_demand(city, job_name, data):
    job_data = data[(data['city'] == city) & (data['positionName'] == job_name)]
    count = len(job_data)
    return {f"{city}{job_name}岗位数量": count}

    # 查询掌握某种技术可以从事的岗位
def query_language_related_jobs(data, language="C++", top_n=5):
    # 动态筛选指定语言的岗位
    jobs = data[data['skillLables'].str.contains(language, na=False, case=False)]
    if jobs.empty:
        return {}
    return jobs['positionName'].value_counts().head(top_n).to_dict()

# === 自定义函数：查询某城市岗位数量最多的五个职位 ===
def query_top_jobs_in_city(city, data, top_n=5):
    """查询指定城市岗位数量最多的前N个职位"""
    city_data = data[data['city'] == city]
    if city_data.empty:
        return "未找到该城市的职位数据"
    return city_data['positionName'].value_counts().head(top_n).to_dict()

    # 查询某种技术的学习推荐
def query_cpp_learning_recommendations(data, top_n=5):
    # 筛选包含C++技能的岗位（假设数据字段为'skillLables'，用逗号分隔技能）
    cpp_jobs = data[data['skillLables'].str.contains('C++', na=False, case=False)]

    if cpp_jobs.empty:
        # 无实际数据时使用示例推荐
        learning_recommendations = {
            "现代C++特性（C++11/14/17/20）": 9,
            "操作系统与系统编程（Linux内核）": 8,
            "高性能计算与并发编程": 7,
            "C++标准库与STL深度": 6,
            "开源项目贡献（如LLVM）": 5
        }
        # 修复括号未闭合问题
        return dict(sorted(learning_recommendations.items(), key=lambda x: x[1], reverse=True)[:top_n])

    # 收集所有关联技能（排除C++本身）
    all_related_skills = []
    for skills in cpp_jobs['skillLables']:
        if isinstance(skills, str):
            skills_list = [s.strip() for s in skills.split(',') if s.strip() != 'C++']
            all_related_skills.extend(skills_list)

    # 统计技能出现频率（作为推荐指数）
    skill_counts = pd.Series(all_related_skills).value_counts().head(top_n)

    # 转换为{技能名称: 推荐指数}格式
    return skill_counts.to_dict()

    # 绘制图表
def plot_top_jobs(top_jobs, chart_type='bar', title=None):
    if not isinstance(top_jobs, dict):
        st.error("图表数据格式错误，无法显示")
        return
    
    options = {
        "title": {"text": title if title else "数据分布"},
        "tooltip": {"trigger": "axis", "axisPointer": {"type": "shadow"}},
        "xAxis": {"type": "category", "data": list(top_jobs.keys()), "axisLabel": {"rotate": 45}},
        "yAxis": {"type": "value", "name": "数量"},
        "series": [{
            "name": "数量",
            "type": chart_type,
            "data": list(top_jobs.values()),
            "itemStyle": {"color": "#5470C6"}
        }]
    }
    
    st_echarts(options=options, height="400px")

# === 增强的AI回复处理逻辑 ===
def handle_tool_calls(tool_calls):
    results = []
    for tool_call in tool_calls:
        func_name = tool_call.function.name
        args = json.loads(tool_call.function.arguments)
        
        if func_name == "query_top_jobs_in_city":
            city = args["city"]
            top_n = args.get("top_n", 5)
            job_counts = query_top_jobs_in_city(city, data, top_n)
            
            # 将结果转换为饼图配置
            if func_name == "query_top_jobs_in_city":
                city = args["city"]
                top_n = args.get("top_n", 5)
                job_counts = query_top_jobs_in_city(city, data, top_n)
            
            # 将结果转换为饼图配置
            if isinstance(job_counts, dict):
                chart_option = {
                    "title": {"text": f"{city}岗位数量最多的{top_n}个职位"},
                    "tooltip": {"trigger": "item", "formatter": "{b}: {c} ({d}%)"},
                    "series": [{
                        "name": "岗位数量",
                        "type": "pie",
                        "radius": "65%",
                        "data": [{"value": v, "name": k} for k, v in job_counts.items()],
                        "emphasis": {"itemStyle": {"shadowBlur": 10, "shadowOffsetX": 0, "shadowColor": "rgba(0, 0, 0, 0.5)"}}
                    }]
                }
                results.append({
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": func_name,
                    "content": json.dumps({
                        "city": city,
                        "job_counts": job_counts,
                        "chart": chart_option
                    })
                })
            else:
                results.append({
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": func_name,
                    "content": json.dumps({"error": job_counts})
                })

        elif func_name == "query_average_salary":
            filtered = data[(data['city'] == args["city"]) & 
                          (data['positionName'] == args["job_name"])]
            avg_salary = filtered['salary_mid'].mean() if not filtered.empty else 0
            results.append({
                "tool_call_id": call.id,  # 修正为 call.id
                "role": "tool",
                "name": func_name,
                "content": json.dumps({"average_salary": avg_salary})
            })
        elif func_name == "query_max_salary_job":
            filtered = data[data['city'] == args["city"]]
            top_jobs = filtered.nlargest(args.get("top_n", 5), 'avg_salary')
            results.append({
                "tool_call_id": call.id,  # 修正为 call.id
                "role": "tool",
                "name": func_name,
                "content": top_jobs.to_json(orient='records')
            })
        elif func_name == "query_salary_range":
            filtered = data[(data['city'] == args["city"]) & 
                          (data['positionName'] == args["job_name"])]
            salary_range = {
                "min": filtered['min_salary'].min(),
                "max": filtered['max_salary'].max()
            }
            results.append({
                "tool_call_id": call.id,  # 修正为 call.id
                "role": "tool",
                "name": func_name,
                "content": json.dumps(salary_range)
            })
        elif func_name == "query_job_demand":
            count = len(data[(data['city'] == args["city"]) & 
                           (data['positionName'] == args["job_name"])])
            results.append({
                "tool_call_id": call.id,  # 修正为 call.id
                "role": "tool",
                "name": func_name,
                "content": json.dumps({"count": count})
            })
        elif func_name == "query_language_related_jobs":
            language = args["language"]
            top_n = args.get("top_n", 5)
            
            # 从数据集中筛选包含该编程语言的岗位（假设技能标签字段为 skillLables）
            # 注意：确保数据集中 skillLables 字段已正确解析为列表（如 ["Python", "Java"]）
            filtered = data[
                data['skillLables'].apply(lambda skills: language in skills)
            ]
            
            # 统计岗位需求量（每个岗位出现的次数）
            job_demand = filtered['positionName'].value_counts().head(top_n)
            # 转换为 DataFrame 格式（用于后续可视化）
            job_demand_df = pd.DataFrame({
                "positionName": job_demand.index.tolist(),
                "demand": job_demand.tolist()
            })
            
            # 结果转换为 JSON 字符串（供 pd.read_json 解析）
            results.append({
                "tool_call_id": call.id,  # 修正为 call.id
                "role": "tool",
                "name": func_name,
                "content": job_demand_df.to_json(orient="records")
            })
        elif func_name == "query_language_learning_recommendations":
            language = args["language"]
            top_n = args.get("top_n", 10)
            # 这里可以添加根据语言返回不同学习推荐的逻辑
            recommendations = {
                "Python": ["基础语法", "Web开发", "数据分析", "机器学习", "自动化脚本"],
                "Java": ["核心语法", "Spring框架", "微服务", "JVM原理", "并发编程"],
                "C++": ["现代C++特性", "操作系统与系统编程", "高性能计算", "C++标准库", "开源项目贡献"],
                "JavaScript": ["ES6+特性", "Node.js", "前端框架", "TypeScript", "Web性能优化"]
            }
            return {
                "tool_call_id": call.id,  # 修正为 call.id
                "role": "tool",
                "name": func_name,
                "content": json.dumps(recommendations.get(language, [])[:top_n])
            }
        elif func_name == "query_cpp_learning_recommendations":
            pass
        elif func_name == "query_work_experience_requirement":
            city = args["city"]
            job_name = args["job_name"]
            # 从数据集中筛选对应城市和职位的记录，统计工作经验分布
            filtered = data[
                (data['city'] == city) & 
                (data['positionName'] == job_name)
            ]
            # 统计每个工作经验的岗位数量
            work_exp_counts = filtered['workYear'].value_counts().reset_index()
            work_exp_counts.columns = ["work_year", "count"]
            # 转换为JSON格式返回（包含城市和职位信息）
            results.append({
                "name": func_name,
                "content": json.dumps({
                    "city": city,
                    "job_name": job_name,
                    "work_experience": work_exp_counts.to_dict("records")
                })
            })
    return results

# === 增强的图表生成函数 ===
def generate_chart(message):
    # 根据AI回复内容生成相应图表
    if "薪资" in message.content:
        return {
            "xAxis": {"type": "category", "data": ["最低", "平均", "最高"]},
            "yAxis": {"type": "value"},
            "series": [{
                "data": [
                    data['min_salary'].mean(),
                    data['avg_salary'].mean(),
                    data['max_salary'].mean()
                ],
                "type": "bar"
            }]
        }
    elif "地区" in message.content:
        region_counts = data['city'].value_counts().head(10)
        return {
            "xAxis": {"type": "category", "data": region_counts.index.tolist()},
            "yAxis": {"type": "value"},
            "series": [{"data": region_counts.tolist(), "type": "bar"}]
        }
    elif "技能" in message.content:
        # 新增技能关联图表
        skill = message.content.split("技能")[1].strip()
        if skill in pd.Series(sum(data['skillLables'].tolist(), [])).unique():
            skill_positions = data[data['skillLables'].apply(lambda x: skill in x)]['positionName'].value_counts()
            return {
                "series": [{
                    "type": "pie",
                    "data": [{"value": v, "name": k} for k, v in skill_positions.items()],
                    "radius": "50%"
                }]
            }

    # 新增编程语言相关岗位图表
    if "编程语言" in message.content or "语言" in message.content:
        language = next((word for word in message.content.split() 
                        if word in ["Python", "Java", "C++", "JavaScript"]), None)
        if language:
            language_jobs = data[data['skillLables'].apply(lambda x: language in x)]
            position_counts = language_jobs['positionName'].value_counts().head(10)
            return {
                "xAxis": {"type": "category", "data": position_counts.index.tolist()},
                "yAxis": {"type": "value"},
                "series": [{"data": position_counts.tolist(), "type": "bar"}],
                "title": {"text": f"{language}相关热门岗位"}
            }

    # 新增最高薪资岗位图表
    if "最高薪资" in message.content or "薪资最高" in message.content:
        city = next((word for word in message.content.split() 
                    if word in data['city'].unique()), None)
        if city:
            top_jobs = data[data['city'] == city].nlargest(5, 'avg_salary')
            return {
                "xAxis": {"type": "category", "data": top_jobs['positionName'].tolist()},
                "yAxis": {"type": "value"},
                "series": [{"data": top_jobs['avg_salary'].tolist(), "type": "bar"}],
                "title": {"text": f"{city}薪资最高岗位"}
            }
    elif "学习推荐" in message.content:
        if "Java" in message.content:
            java_jobs = data[data['skillLables'].apply(lambda x: "Java" in x)]
            all_related_skills = []
            for skills in java_jobs['skillLables']:
                if isinstance(skills, list):
                    skills_list = [s.strip() for s in skills if s.strip() != "Java"]
                    all_related_skills.extend(skills_list)
            skill_counts = pd.Series(all_related_skills).value_counts().head(5)
            return {
                "series": [{
                    "type": "pie",
                    "radius": ['30%', '70%'],
                    "data": [{"value": v, "name": k} for k, v in skill_counts.items()]
                }]
            }
        elif "C++" in message.content:
            cpp_jobs = data[data['skillLables'].apply(lambda x: "C++" in x)]
            all_related_skills = []
            for skills in cpp_jobs['skillLables']:
                if isinstance(skills, list):
                    skills_list = [s.strip() for s in skills if s.strip() != "C++"]
                    all_related_skills.extend(skills_list)
            skill_counts = pd.Series(all_related_skills).value_counts().head(5)
            return {
                "series": [{
                    "type": "pie",
                    "radius": ['30%', '70%'],
                    "data": [{"value": v, "name": k} for k, v in skill_counts.items()]
                }]
            }
        # 新增工作经验要求图表
    if tool_results and "query_experience_requirement" in message.content:
        work_exp_data = json.loads(tool_results[0]['content'])
        chart_option = {
            "title": {"text": f"{work_exp_data['city']}{work_exp_data['job_name']}工作经验要求分布"},
            "tooltip": {
                "trigger": "axis",
                "formatter": "{b}: {c} 个岗位"
            },
            "xAxis": {
                "type": "category",
                "data": [item['work_year'] for item in work_exp_data['work_experience']]
            },
            "yAxis": {
                "type": "value",
                "name": "岗位数量"
            },
            "series": [{
                "type": "bar",
                "data": [item['count'] for item in work_exp_data['work_experience']],
                "itemStyle": {
                    "color": "#5470C6",
                    "barBorderRadius": 5
                },
                "label": {"show": True, "formatter": "{c}"}
            }]
        }
        return chart_option
    
    return None


def format_tool_response(tool_results):
    """格式化工具调用结果为自然语言回答"""
    responses = []
    for result in tool_results:
        func_name = result.get("name")
        content = json.loads(result.get("content", "{}"))
        
        if func_name == "query_average_salary":
            avg_salary = content.get("average_salary", 0)
            responses.append(f"该职位的平均薪资为：{avg_salary:.2f}元")
        
        elif func_name == "query_job_demand":
            count = content.get("count", 0)
            responses.append(f"该城市的岗位需求量为：{count}个")
        
        elif func_name == "query_salary_range":
            min_salary = content.get("min", 0)
            max_salary = content.get("max", 0)
            responses.append(f"薪资区间为：{min_salary:.2f}元 - {max_salary:.2f}元")
        
        elif func_name in ["query_language_jobs", "query_max_salary_job"]:
            jobs = pd.read_json(content)
            if not jobs.empty:
                responses.append(f"查询到以下岗位信息：\n{jobs[['positionName', 'avg_salary']].to_string(index=False)}")
        
        elif func_name == "query_language_related_jobs":
            position_counts = pd.read_json(content)
            responses.append(f"掌握该语言可从事的热门岗位及其需求量：\n{position_counts.to_string()}")
        
        elif func_name == "query_language_learning_recommendations":
            recommendations = content
            responses.append(f"学习推荐：{', '.join(recommendations)}")

        if func_name == "query_top_jobs_in_city":
            if "error" in content:
                responses.append(f"{content['error']}")
            else:
                city = content.get("city")
                job_counts = content.get("job_counts")
                chart_option = content.get("chart")
                
                responses.append(f"{city}岗位数量最多的五个职位为：")
                for job, count in job_counts.items():
                    responses.append(f"- {job}: {count}个岗位")
                
                # 保存图表数据以便在界面中显示
                if chart_option:    
                    st.session_state.last_chart_data = chart_option

    return "\n\n".join(responses) if responses else "未获取到有效数据"


# === 处理用户问题的核心函数 ===
def process_user_question(question, data, api_key):
    try:
        client = init_ai(api_key)
        messages = [{"role": "user", "content": question}]

        if st.session_state.messages:
            messages.extend([msg for msg in st.session_state.messages[-5:]])

        response = client.chat.completions.create(
            model="glm-4",
            messages=messages,
            tools=tools,
            tool_choice="auto"
        )

        if response.choices[0].message.tool_calls:
            tool_call = response.choices[0].message.tool_calls[0]
            func_name = tool_call.function.name
            params = json.loads(tool_call.function.arguments)

            if func_name == "query_top_jobs_in_city":
                city = params["city"]
                job_counts = query_top_jobs_in_city(city, data)
                
                if isinstance(job_counts, dict):
                    # 生成自然语言回答
                    answer = f"{city}岗位数量最多的五个职位为：\n"
                    for job, count in job_counts.items():
                        answer += f"- {job}: {count}个岗位\n"
                    
                    # 生成饼图
                    chart_option = {
                        "title": {"text": f"{city}岗位数量最多的五个职位"},
                        "tooltip": {"trigger": "item", "formatter": "{b}: {c} ({d}%)"},
                        "series": [{
                            "name": "岗位数量",
                            "type": "pie",
                            "radius": "65%",
                            "data": [{"value": v, "name": k} for k, v in job_counts.items()],
                            "emphasis": {"itemStyle": {"shadowBlur": 10, "shadowOffsetX": 0, "shadowColor": "rgba(0, 0, 0, 0.5)"}}
                        }]
                    }
                    
                    st.session_state.last_chart_data = chart_option
                    return answer
                else:
                    return job_counts
                    
            elif func_name == "query_average_salary":
                result = query_average_salary(params["city"], params["job_name"], data)
                chart_type = 'bar'
                title = f"{params['city']}{params['job_name']}平均薪资"
                # 确保result是字典格式
                if isinstance(result, (float, int)):
                    result = {title: result}

            elif func_name == "query_salary_range":
                result = query_salary_range(params["city"], params["job_name"], data)

            elif func_name == "query_experience_requirement":
                result = query_experience_requirement(params["city"], params["job_name"])
                
                # 生成自然语言回答
                prompt = f"""根据以下信息回答问题：
                用户提问：{question}
                查询结果：{result}
                请用简洁自然的语言回答，不超过200字"""
                
                final_response = client.chat.completions.create(
                    model="glm-4",
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.6
                )
                answer = final_response.choices[0].message.content
                
                # 生成图表
                chart_option = generate_chart(answer, tool_results=[{
                    "name": func_name,
                    "content": json.dumps({
                        "city": params["city"],
                        "job_name": params["job_name"],
                        "work_experience": result
                    })
                }])
                
                if chart_option:
                    st.session_state.last_chart_data = chart_option
                return answer

            elif func_name == "query_job_demand":
                result = query_job_demand(params["city"], params["job_name"], data)
                # 确保result是字典格式
                if isinstance(result, int):
                    result = {f"{params['city']}{params['job_name']}岗位数量": result}
                chart_type = 'bar'
                title = f"{params['city']}{params['job_name']}岗位数量"

            elif func_name == "query_max_salary_job":
                result = query_max_salary_job(params["city"], data, top_n=1)
                chart_type = 'bar'  # 修改为柱状图
                title = f"{params['city']}薪资最高岗位"  # 添加标题

            elif func_name == "query_language_related_jobs":
                language = params["language"]
                result = query_language_related_jobs(data, language=language, top_n=params.get("top_n", 5))
                chart_type = 'pie'  # 新增：显式设置为饼图
                title = f"{language}相关热门岗位分布"  # 可选：添加标题（若需要）

            # 新增：处理学习推荐工具
            elif func_name == "query_cpp_learning_recommendations":
                result = query_cpp_learning_recommendations(data, params.get("top_n", 10))
                chart_type = 'ring'  # 关键修改：将饼图改为环形图
                title = None  # 保持标题不显示

            # 生成自然语言回答
            prompt = f"""根据以下信息回答问题：
            用户提问：{question}
            查询结果：{result}
            请用简洁自然的语言回答，不超过200字"""

            final_response = client.chat.completions.create(
                model="glm-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.6
            )
            answer = final_response.choices[0].message.content

            # 关键修改：调整保存条件（允许title为None）
            if chart_type and isinstance(result, dict):  # 加上这个判断
                st.session_state.last_chart_data = (result, chart_type, title)

            return answer
        else:
            return response.choices[0].message.content
    except Exception as e:
        return f"处理问题时发生错误: {str(e)}"

        # 主函数
def main():

    st.title("招聘数据智能助手")

    # 统一初始化所有会话状态变量
    if 'messages' not in st.session_state:
        st.session_state.messages = [{"role": "assistant", "content": "您好，我是招聘数据分析助手，请问有什么可以帮助您？"}]
    if 'api_key' not in st.session_state:
        st.session_state.api_key = ""
    if 'last_chart_data' not in st.session_state:  # 改为存储数据而非widget
        st.session_state.last_chart_data = None  # 初始化图表数据状态
    if 'last_table_data' not in st.session_state:  # 新增：初始化表格数据状态
        st.session_state.last_table_data = None  # 初始化为None
    if 'logged_in' not in st.session_state:  # 新增登录状态初始化
        st.session_state.logged_in = False  # 默认未登录

    data = load_data()

    # 左侧边栏内容
    with st.sidebar:
        st.header("功能选项")
        # 新增登录表单
        st.subheader("用户登录")
        username = st.text_input("用户名", key="login_username")
        password = st.text_input("密码", type="password", key="login_password")
        if st.button("登录"):
            if username == "admin" and password == "123456":
                st.session_state.logged_in = True
                st.success("登录成功！")
            else:
                st.error("用户名或密码错误")
        api_key = st.text_input("请输入您的API Key", type="password", value=st.session_state.api_key, help="请输入有效的智谱AI API Key")
        if api_key != st.session_state.api_key:
            st.session_state.api_key = api_key
            st.rerun()
        def clear_session():
            st.session_state.messages = [{"role": "assistant", "content": "您好，我是招聘数据分析助手，请问有什么可以帮助您？"}]
            st.session_state.last_chart_data = None
            st.session_state.last_table_data = None

        st.button('清空聊天记录', on_click=clear_session)

    # 显示历史对话
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])

    # 显示图表
    if st.session_state.last_chart_data:
        if isinstance(st.session_state.last_chart_data, dict):
            st_echarts(options=st.session_state.last_chart_data, height="400px")
        elif isinstance(st.session_state.last_chart_data, tuple) and len(st.session_state.last_chart_data) == 3:
            result, chart_type, title = st.session_state.last_chart_data
            if isinstance(result, dict):
                plot_top_jobs(result, chart_type=chart_type, title=title)
            else:
                st.warning("图表数据格式错误，无法显示")
        else:
            st.warning("图表数据格式错误，无法显示")

    # 在主函数main中，展示表格的部分
    if st.session_state.last_table_data is not None:
        st.markdown("### 📋 岗位详细信息")
        st.dataframe(
            st.session_state.last_table_data,
            height=500,
            use_container_width=True,
            hide_index=True,
            column_config={
                "职位名称": "职位名称",
                "公司名称": st.column_config.TextColumn(
                    "公司名称", width="large"
                ),
                "城市": st.column_config.TextColumn(
                    "城市", width="small"
                ),
                "薪资范围": st.column_config.ProgressColumn(
                    "薪资水平",
                    help="根据薪资中位数计算的相对水平",
                    format="%.2f",
                    min_value=0,
                    max_value=100
                )
            }
        )
    elif st.session_state.last_table_data is not None and st.session_state.last_table_data.empty:
        st.warning("未找到相关岗位信息")

    # 只有输入了有效的API key且已登录才能使用核心功能
    if st.session_state.api_key and st.session_state.logged_in:
        try:
            client = init_ai(st.session_state.api_key)
            if prompt := st.chat_input("请输入您的问题"):
                # 添加用户消息
                st.session_state.messages.append({"role": "user", "content": prompt})
                with st.chat_message("user"):
                    st.write(prompt)

                # 处理问题
                with st.spinner("正在思考..."):
                    with st.chat_message("assistant"):
                        answer = process_user_question(prompt, data, st.session_state.api_key)
                        if answer:
                            st.write(answer)
                    st.session_state.messages.append({"role": "assistant", "content": answer})

                # 关键修改：在问题处理完成后显示图表（确保last_chart_data已更新）
                if st.session_state.last_chart_data:
                    result, chart_type, title = st.session_state.last_chart_data
                    plot_top_jobs(result, chart_type=chart_type, title=title)
        except Exception as e:
            st.error(f"发生错误: {str(e)}")
            st.session_state.api_key = ""

    elif not st.session_state.logged_in:  # 未登录提示
            st.warning("请先在左侧边栏完成登录以使用AI咨询功能")
    else:
        st.warning("请输入有效的API Key以启用聊天功能")

if __name__ == "__main__":
    main()